# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import functools
from contextlib import contextmanager

import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

from coin.base.datetime_util import to_datetime


def setup_plt(figsize_y=None):
  plt.style.use('fivethirtyeight')
  plt.rcParams['lines.linewidth'] = 0.5
  plt.rcParams['figure.dpi'] = 200
  plt.rcParams['figure.figsize'] = [8, figsize_y or 3]
  plt.rcParams['font.family'] = ['Inconsolata']
  plt.rcParams['font.sans-serif'] = ['Inconsolata']
  plt.rcParams['font.monospace'] = ['Inconsolata']
  plt.rcParams['font.size'] = 6
  plt.rcParams['legend.fontsize'] = 6
  plt.rcParams['xtick.labelsize'] = 6
  plt.rcParams['ytick.labelsize'] = 6
  plt.rcParams['savefig.facecolor'] = '#ffffff'


def _plot(names, ranges, ax, name, times, values):
  assert len(times) == len(values)
  if not times:
    return
  if name is not None:
    names.append(name)
  ranges.append((min(times), max(times)))
  df = pd.DataFrame({'time': times, 'value': values})
  ax.plot(df['time'], df['value'])


@contextmanager
def get_plotter(filepath,
                begin_time=None,
                end_time=None,
                plot_secondary=False,
                axis_interval_hour=2):
  names = []
  ranges = []
  setup_plt(4.5 if plot_secondary else 3)

  if plot_secondary:
    fig = plt.figure()
    ax1 = plt.subplot2grid((3, 1), (0, 0), rowspan=2)
    ax2 = plt.subplot2grid((3, 1), (2, 0), sharex=ax1)
    axes = [ax1, ax2]
  else:
    fig, ax1 = plt.subplots()
    axes = [ax1]

  # fig.patch.set_facecolor('#ffffff')
  plotters = [functools.partial(_plot, names, ranges, ax) for ax in axes]
  yield plotters if plot_secondary else plotters[0]

  if not names:
    plt.close(fig)
    return

  if begin_time is not None:
    begin_time = min([r[0] for r in ranges])
  begin_time = to_datetime(begin_time)

  if end_time is None:
    end_time = max([r[1] for r in ranges])
  end_time = to_datetime(end_time)

  if axis_interval_hour is not None:
    ax1.set_xlim((begin_time, end_time))
    ax1.xaxis.set_major_locator(mdates.HourLocator(interval=axis_interval_hour))
    ax1.xaxis.set_major_formatter(mdates.DateFormatter('%m-%d %02HH'))
    fig.autofmt_xdate()

  ax1.legend(names, loc='upper left')
  fig.tight_layout()
  fig.savefig(filepath, transparent=False)
  plt.close(fig)


def refine_balance(result, begin_time, end_time, mark_time):
  begin_time = to_datetime(begin_time)
  mark_time = to_datetime(mark_time or begin_time)
  end_time = to_datetime(end_time)

  name = result['name']
  times = []
  values = []

  base_value = None
  for elem in result['balance_dump']:
    if elem['time'] < mark_time:
      base_value = elem['value']
    else:
      if base_value is None:
        base_value = elem['value']
      break

  for elem in result['balance_dump']:
    if begin_time <= elem['time'] <= end_time:
      times.append(elem['time'])
      values.append(elem['value'] - base_value)

  return name, times, values


def refine_timeseries(result, begin_time, end_time, key, emit_name=False):
  begin_time = to_datetime(begin_time)
  end_time = to_datetime(end_time)
  times = []
  values = []

  for elem in result[key]:
    if begin_time <= elem['time'] <= end_time:
      times.append(elem['time'])
      values.append(elem['value'])

  return result['name'] if emit_name else None, times, values
